#!/usr/bin/env python3
# -*- coding: utf-8
"""
知识库服务
使用Weaviate存储和检索PDF文档内容，支持增强检索
"""

import uuid
from typing import List, Dict, Any, Optional
import weaviate
from weaviate.classes.config import Property, DataType, Configure
from config.weaviate_config import get_weaviate_client_config, get_openai_embedding_config
from utils.logger import emobot_logger

logger = emobot_logger.get_logger()


class KnowledgeBaseService:
    """知识库服务类"""
    
    def __init__(self):
        """初始化知识库服务"""
        try:
            # 获取Weaviate配置
            weaviate_config = get_weaviate_client_config()
            embedding_config = get_openai_embedding_config()
            
            # 创建Weaviate客户端
            self.client = weaviate.connect_to_local(
                host=weaviate_config["weaviate_url"].replace("http://", "").split(":")[0],
                port=int(weaviate_config["weaviate_url"].split(":")[-1]),
                grpc_port=50051
            )
            
            self.collection_name = "PsychologyKnowledge"
            self.embedding_config = embedding_config
            
            logger.info(f"知识库服务初始化完成 - Weaviate: {weaviate_config['weaviate_url']}")
            
        except Exception as e:
            logger.error(f"知识库服务初始化失败: {e}")
            raise
    
    def create_collection(self) -> bool:
        """
        创建心理学知识库集合
        
        Returns:
            bool: 创建是否成功
        """
        try:
            # 检查集合是否已存在
            if self.client.collections.exists(self.collection_name):
                logger.info(f"集合 {self.collection_name} 已存在")
                return True
            
            # 创建集合
            self.client.collections.create(
                name=self.collection_name,
                properties=[
                    Property(name="content", data_type=DataType.TEXT, description="文本内容"),
                    Property(name="file_name", data_type=DataType.TEXT, description="文件名"),
                    Property(name="file_path", data_type=DataType.TEXT, description="文件路径"),
                    Property(name="page_number", data_type=DataType.INT, description="页码"),
                    Property(name="chunk_index", data_type=DataType.INT, description="块索引"),
                    Property(name="content_type", data_type=DataType.TEXT, description="内容类型"),
                    Property(name="char_count", data_type=DataType.INT, description="字符数"),
                    Property(name="metadata", data_type=DataType.TEXT, description="元数据JSON")
                ]
            )
            
            logger.info(f"集合 {self.collection_name} 创建成功")
            return True
            
        except Exception as e:
            logger.error(f"创建集合失败: {e}")
            return False
    
    def store_pdf_chunks(self, chunks: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        将PDF文本块存储到Weaviate
        
        Args:
            chunks: PDF文本块列表
            
        Returns:
            存储结果
        """
        try:
            # 确保集合存在
            if not self.create_collection():
                return {"success": False, "error": "创建集合失败"}
            
            collection = self.client.collections.get(self.collection_name)
            stored_count = 0
            failed_count = 0
            
            for chunk in chunks:
                try:
                    # 准备存储数据
                    data_object = {
                        "content": chunk["text"],
                        "file_name": chunk.get("file_name", ""),
                        "file_path": chunk.get("file_path", ""),
                        "page_number": chunk.get("page_number", 1),
                        "chunk_index": chunk.get("chunk_index", 0),
                        "content_type": chunk.get("content_type", "psychology_pdf"),
                        "char_count": chunk.get("char_count", 0),
                        "metadata": str(chunk.get("metadata", {}))
                    }
                    
                    # 存储到Weaviate
                    collection.data.insert(data_object)
                    stored_count += 1
                    
                except Exception as e:
                    logger.error(f"存储文本块失败: {e}")
                    failed_count += 1
            
            result = {
                "success": True,
                "stored_count": stored_count,
                "failed_count": failed_count,
                "total_chunks": len(chunks)
            }
            
            logger.info(f"PDF文本块存储完成: 成功 {stored_count}, 失败 {failed_count}")
            return result
            
        except Exception as e:
            logger.error(f"存储PDF文本块失败: {e}")
            return {"success": False, "error": str(e)}
    
    def search_similar_content(self, query: str, limit: int = 3, similarity_threshold: float = 0.7) -> List[Dict[str, Any]]:
        """
        搜索相似内容
        
        Args:
            query: 查询文本
            limit: 返回结果数量限制
            similarity_threshold: 相似度阈值
            
        Returns:
            相似内容列表
        """
        try:
            collection = self.client.collections.get(self.collection_name)
            
            # 执行BM25文本搜索（不依赖向量化器）
            response = collection.query.bm25(
                query=query,
                limit=limit,
                return_metadata=["score"],
                return_properties=["content", "file_name", "page_number", "chunk_index", "char_count"]
            )
            
            results = []
            for obj in response.objects:
                result = {
                    "content": obj.properties.get("content", ""),
                    "file_name": obj.properties.get("file_name", ""),
                    "page_number": obj.properties.get("page_number", 1),
                    "chunk_index": obj.properties.get("chunk_index", 0),
                    "char_count": obj.properties.get("char_count", 0),
                    "similarity_score": 0.8,  # BM25搜索使用固定相似度分数
                    "bm25_score": obj.metadata.score if hasattr(obj.metadata, 'score') else 0.8
                }
                results.append(result)
            
            logger.info(f"搜索完成: 查询 '{query[:50]}...', 返回 {len(results)} 个结果")
            return results
            
        except Exception as e:
            logger.error(f"搜索相似内容失败: {e}")
            return []
    
    def get_collection_stats(self) -> Dict[str, Any]:
        """
        获取集合统计信息
        
        Returns:
            统计信息
        """
        try:
            collection = self.client.collections.get(self.collection_name)
            
            # 获取对象总数
            total_objects = collection.aggregate.over_all(total_count=True).total_count
            
            # 获取文件统计
            response = collection.query.fetch_objects(
                return_properties=["file_name"],
                limit=1000  # 获取足够多的对象来统计
            )
            
            file_stats = {}
            for obj in response.objects:
                file_name = obj.properties.get("file_name", "unknown")
                file_stats[file_name] = file_stats.get(file_name, 0) + 1
            
            stats = {
                "total_objects": total_objects,
                "unique_files": len(file_stats),
                "file_distribution": file_stats,
                "collection_name": self.collection_name
            }
            
            logger.info(f"集合统计信息: {stats}")
            return stats
            
        except Exception as e:
            logger.error(f"获取集合统计信息失败: {e}")
            return {"error": str(e)}
    
    def clear_collection(self) -> bool:
        """
        清空集合
        
        Returns:
            bool: 清空是否成功
        """
        try:
            collection = self.client.collections.get(self.collection_name)
            
            # 删除所有对象
            collection.data.delete_many()
            
            logger.info(f"集合 {self.collection_name} 已清空")
            return True
            
        except Exception as e:
            logger.error(f"清空集合失败: {e}")
            return False
    
    def close(self):
        """关闭Weaviate连接"""
        try:
            if hasattr(self, 'client'):
                self.client.close()
                logger.info("Weaviate连接已关闭")
        except Exception as e:
            logger.error(f"关闭Weaviate连接失败: {e}")


def create_knowledge_base_service() -> KnowledgeBaseService:
    """
    创建知识库服务实例
    
    Returns:
        KnowledgeBaseService实例
    """
    return KnowledgeBaseService()
